#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
作  者    : 北极星光 light22@126.com
创建时间  : 2023-07-17 00:29:37
最后编辑  : 2023-08-26 17:33:37
编辑人    : 北极星光
'''

# apiKey 获取地址： https://console.bce.baidu.com/iam/#/iam/apikey/list
# 支持的模型列表： https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Fm2vrveyu
# 文档地址： https://cloud.baidu.com/doc/qianfan-api/s/3m7of64lb

import yaml
import requests
import json
from datetime import datetime

class ErnieAI:

    def __init__(self,
                 api_key=None,
                 bot_name=None,
                 bot_role=None,
                 welcome_words=None,
                 sensitive_words=None):
        self.api_key = api_key
        self.bot_name = bot_name
        self.bot_role = bot_role
        self.welcome_words = welcome_words
        self.sensitive_words = sensitive_words
        self.messages = {
            'default_user': [{
                'role': 'system',
                'content': self.bot_role
            }]
        }

    # 文本生成
    def ernie_text_generate(self, text, model_name, user='default_user'):
        if user not in self.messages:
            self.messages[user] = [{
                'role': 'system',
                'content': self.bot_role
            }]
        self.messages[user].append({'role': 'user', 'content': text})

        # 过滤敏感词
        if any(word in text for word in self.sensitive_words):
            response = '''### 文心一言 AI文本生成
            消息中含有敏感词，请重新输入。'''
            return response

        # 文本生成
        url = "https://qianfan.baidubce.com/v2/chat/completions"

        payload = json.dumps({
            "model": model_name,
            "messages": self.messages[user],
        })
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        response = requests.request("POST", url, headers=headers, data=payload)
        self.messages[user].append(
            response.json().get('choices')[0].get('message'))

        return response.json().get('choices')[0].get('message').get('content')
        # return response.text
    
    # 图像生成
    def ernie_image_generate(self, prompt, model_name):
        url = "https://qianfan.baidubce.com/v2/images/generations"

        payload = json.dumps({
            "model": model_name,
            "prompt": prompt
        })
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        response = requests.request("POST", url, headers=headers, data=payload)

        return response.json()

    # 获取模型列表
    def list_models(self):
        url = "https://qianfan.baidubce.com/v2/models"

        payload = json.dumps({})
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }

        response = requests.request("GET", url, headers=headers, data=payload)
        model_list = [i.get('id') for i in response.json().get('data')]
        model_list_text = f'''### 文心一言 所有模型列表'''
        for model in model_list:
            model_list_text += f'''
- {model}'''

        return model_list_text
    
    # 查询余额（暂时不可用，以后有空研究）
    # def get_balance(self):
    #     url = "http://billing.baidubce.com/v1/finance/cash/balance"
    #     http_method = "POST"
    #     headers = {
    #         "Host": "billing.baidubce.com",
    #         "x-bce-date": datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ")  # 使用百度要求的日期格式
    #     }

    #     # 生成签名逻辑 (简化版，完整流程需参考百度官方文档)
    #     canonical_uri = "/v1/finance/cash/balance"
    #     canonical_headers = f"host:{headers['Host']}\nx-bce-date:{headers['x-bce-date']}\n"
    #     signed_headers = "host;x-bce-date"
    #     canonical_request = f"{http_method}\n{canonical_uri}\n\n{canonical_headers}\n{signed_headers}\nUNSIGNED-PAYLOAD"

    #     signing_key = hmac.new(SECRET_KEY.encode(), "bce-auth-v1/{}/{}/1800".format(ACCESS_KEY, headers['x-bce-date']).encode(), hashlib.sha256).hexdigest()
    #     signature = hmac.new(signing_key.encode(), canonical_request.encode(), hashlib.sha256).hexdigest()

    #     # 构造 Authorization 头
    #     headers["Authorization"] = f"bce-auth-v1/{ACCESS_KEY}/{headers['x-bce-date']}/1800/host;x-bce-date/{signature}"

    #     response = requests.post(url, headers=headers)
        
    #     return response.text

    # 重置消息记录
    def reset_messages(self, user='default_user'):
        self.messages[user] = [{
            'role': 'system',
            'content': self.bot_role
        }]


if __name__ == '__main__':
    with open('./config/ernie_config.yaml', 'r', encoding='utf-8') as f:
        cfg = yaml.load(f, Loader=yaml.FullLoader)
    api_key = cfg.get('api_key')
    bot_name = cfg.get('bot_name')
    bot_role = cfg.get('bot_role')
    welcome_words = cfg.get('welcome_words')
    model_name_text_generate = cfg.get('model_name_text_generate')
    model_name_image_generate = cfg.get('model_name_image_generate')
    sensitive_words = cfg.get('sensitive_words')
    ernie_ai = ErnieAI(api_key, bot_name, bot_role, welcome_words,
                       sensitive_words)
    # response = ernie_ai.ernie_text_generate("用最肉麻的话向我求爱", model_name_text_generate)
    # response = ernie_ai.list_models()
    # response = ernie_ai.get_balance()
    response = ernie_ai.ernie_image_generate("", model_name_image_generate)
    print(response)
    print(ernie_ai.messages)
